#!/usr/bin/python3
# Date: Fri Nov 19 16:27:47 2021
# Author: January

from ctypes import pointer
import time
from Crypto.Math.Numbers import Integer
import ecpy
from ecpy.curves import Curve,Point
from Crypto.Cipher import AES
from Crypto.Protocol.KDF import PBKDF2
from Crypto.PublicKey import ECC
import hmac
import hashlib
import random
import socket
import argparse
import logging
import time

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("dronesec")


def rand_bytes(length):
    return b''.join([int.to_bytes(random.randint(0,255), 1, 'little') for i in range(length)])

# 预定义参数
curve = Curve.get_curve('secp256r1')
private_keys = {
    0 : 0xfb26a4e75eec75544c0f44e937dcf5ee6355c7176600b9688c667e5c283b43c5,
    1 : 0xfb26a4e75abcd5544c0f44e937dcf5ee6355c7176600b9688c667e5c283b43c5,
    2 : 0xfb26a4e7523455544c0f44e937dcf5ee6355c7176600b9688c667e5c283b43c5
}

# using pycryptodome ecc(默认曲线为secp256r1)
G = ECC.EccPoint(curve.generator.x, curve.generator.y)

# using ecpy ecc
# G = curve.generator

public_keys = {
    0 : private_keys[0] * G,
    1 : private_keys[1] * G,
    2 : private_keys[2] * G
}

def int_to_bytes(num):
    return int.to_bytes(num, 32, "big")

# ecpy和pycryptodome中ecc运算的性能对比
def point_arithmetic_performance_test():
    curve = Curve.get_curve('secp256r1')
    G = ECC.EccPoint(curve.generator.x, curve.generator.y)
    start = time.time()
    P_crypto = private_keys[0] * G
    end = time.time()
    print(P_crypto.x)
    print(P_crypto.y)
    print("pycryotodome ecc mul time used %.2fms"%((end - start) * 1000))

    start = time.time()
    P_ecpy = private_keys[0] * curve.generator
    end = time.time()
    print(P_ecpy.x)
    print(P_ecpy.y)
    print("ecpy ecc mul time used %.2fms"%((end - start) * 1000))


def performance_test():
    curve = Curve.get_curve('secp256r1')
    G = ECC.EccPoint(curve.generator.x, curve.generator.y)

    a = Integer.from_bytes(rand_bytes(32))
    b = Integer.from_bytes(rand_bytes(32))

    runtimes = 10000

    start = time.time()
    for i in range(runtimes):
       z = a * b
    end = time.time()
    print("mul used %.2fms"%((end - start) * 1000)) 

    start = time.time()
    for i in range(runtimes):
       z = a + b
    end = time.time()
    print("add used %.2fms"%((end - start) * 1000)) 

    field = curve.field
    start = time.time()
    for i in range(runtimes):
        a.inverse(field)
    end = time.time()
    print("inv used %.2fms"%((end - start) * 1000))

    start = time.time()
    for i in range(runtimes):
        P_crypto = private_keys[0] * G
    end = time.time()
    print("ecc mul time used %.2fms"%((end - start) * 1000))

    start = time.time()
    for i in range(runtimes):
        P_crypto = P_crypto + G
    end = time.time()
    print("ecc add time used %.2fms"%((end - start) * 1000))

    data = b'\x01' * 256
    key = b'\x02' * 16
    start = time.time()
    for i in range(runtimes):
        hmac.new(key, data, digestmod=hashlib.sha256).digest()
    end = time.time()
    print("hash time used %.2fms"%((end - start) * 1000)) 

    start = time.time()
    for i in range(runtimes):
        AES.new(key, AES.MODE_ECB).encrypt(data)
    end = time.time()
    print("enc time used %.2fms"%((end - start) * 1000)) 

addrs = {
    0 : ("192.168.3.186", 9000),
    1 : ("192.168.3.108", 9001),
    2 : ("127.0.0.1", 9002)
}

keys = {
    0 : b'\xaf\x99\x85\x99\x05\xea}W\x85\xcf\x0bwA\x1f\xc0\x97',
    1 : b'\xaf\x99\x85\x99\x05\xea}W\x85\xcf\x0bwA\x1f\xc0\x97',
    2 : b'{\xbe0g\xc1$\xfaa\xda\xca\x14^\xbd;Y\xee'
}




def calculate_num_byte_size(num):
    size = 1
    num >>= 8
    while num != 0:
        num >>= 8
        size += 1
    return size

# 将多个字段封装为一个数据包
def pack_bytes(bytes_list):
    packed_bytes = b""
    for item in bytes_list:
        item_len = len(item)
        if item_len <= 128:
            size_field = item_len.to_bytes(1, "little")
        else:
            item_len_bytes_size = calculate_num_byte_size(item_len)
            len_byte = (128 + item_len_bytes_size).to_bytes(1, "little")
            len_bytes = item_len.to_bytes(item_len_bytes_size, "little")
            size_field = len_byte + len_bytes
        packed_bytes = packed_bytes + size_field + item
    return packed_bytes

# 从封装的数据包中解析多个字段
def unpack_bytes(packed_bytes):
    bytes_list = []
    p = 0
    while p < len(packed_bytes):
        field_len = packed_bytes[p]
        p += 1
        if field_len <= 128:
            item_end = p + field_len
            item = packed_bytes[p : item_end]
        else:
            len_bytes_size = field_len - 128
            len_bytes = packed_bytes[p : p + len_bytes_size]
            p += len_bytes_size
            field_len = int.from_bytes(len_bytes, byteorder="little")
            item_end = p + field_len
            item = packed_bytes[p : item_end]
        bytes_list.append(item)
        p = item_end
    return bytes_list

# basic implementation of dronesec protocol
class Dronesec:
    def __init__(self, id : int, curve : Curve, sk, public_keys, addrs):
        self.id = id
        self.public_keys = public_keys
        self.addrs = addrs
        # 长期密钥（预设参数）
        self.curve = curve
        self.sk = sk
        self.channel = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.channel.settimeout(4)
        self.channel.bind(addrs[id])
    
    def accept_key(self, station_id):
        id = int.to_bytes(self.id, 1, "little")
        P_g = self.get_pk(station_id)
        station_id = int.to_bytes(station_id, 1, "little")
        logger.info("Protocol [dronesec], waiting for key...")
        self.channel.settimeout(600)
        msg_g1, addr = self.channel.recvfrom(1024)
        self.channel.settimeout(4)
        start_time = time.time()

        id_g, R_s, mac_g1 = unpack_bytes(msg_g1)
        R_s = self.bytes_to_point(R_s)
        point_vk = R_s + self.sk * P_g
        vk = hashlib.md5(self.point_to_bytes(point_vk)).digest()
        #logger.debug("vk is [%s]"%(vk.hex()))
        if station_id != id_g:
            logger.error("msg from unknown peer")
            return

        if hmac.new(vk, id_g + self.point_to_bytes(R_s), digestmod="sha256").digest() != mac_g1:
            logger.error("msg with invalid mac_g1")
            return
        
        r_u = generate_random(self.curve.field)
        R_u = r_u * G
        point_tmpkey = r_u * R_s
        tmpkey = hashlib.md5(self.point_to_bytes(point_tmpkey)).digest()
        #logger.debug("tmpkey is [%s]"%(tmpkey.hex()))
        mac_u1 = hmac.new(vk, id + self.point_to_bytes(R_u), digestmod="sha256").digest()
        msg_u1 = pack_bytes([id, self.point_to_bytes(R_u), mac_u1])
        self.channel.sendto(msg_u1, addr)

        msg_g2, addr = self.channel.recvfrom(1024)
        id_g2, ekey, ktype, mac_g2 = unpack_bytes(msg_g2)
        if id_g2 != id_g:
            logger.error("msg from unknown peer")
            return
        
        if hmac.new(tmpkey, id_g2 + ekey + ktype, digestmod="sha256").digest() != mac_g2:
            logger.error("msg with invalid mac_g2")
            return
        
        key = AES.new(tmpkey, AES.MODE_ECB).decrypt(ekey)
        mac_u2 = hmac.new(tmpkey, id + key, digestmod="sha256").digest()
        msg_u2 = pack_bytes([id, mac_u2])
        self.channel.sendto(msg_u2, addr)
        ktype = int.from_bytes(ktype, "little")
        logger.info("success to get key [%s], ktype [%d]"%(key.hex(), ktype))
        end_time = time.time()
        logger.info("Protocol finished in [%.2f]ms"%((end_time - start_time) * 1000))
        return key, ktype

    def send_key(self, peer_id, key, ktype):
        # 读取无人机的公钥
        P_u = self.get_pk(peer_id)
        # 读取无人机的ip地址和端口号
        peer_addr = self.get_addr(peer_id)
        #logger.debug("start to set key:\nprotocol [dronesec]\nid [%d]\nPk [%s]\nkey[%s]\nktype[%s]"%(peer_id, self.point_to_bytes(P_u).hex(), key.hex(), ktype) )
        peer_id = int.to_bytes(peer_id, 1, "little")
        id = int.to_bytes(self.id, 1, "little")
        ktype = int.to_bytes(ktype, 1, "little")

        start_time = time.time()

        r_s = generate_random(self.curve.field)
        R_s = r_s * G
        point_vk = R_s + self.sk * P_u
        vk = hashlib.md5(self.point_to_bytes(point_vk)).digest()
        #logger.debug("vk is [%s]"%(vk.hex()))
        mac_g = hmac.new(vk, id + self.point_to_bytes(R_s), digestmod="sha256").digest()
        # 封装第一条消息
        msg_g1 = pack_bytes([id, self.point_to_bytes(R_s), mac_g])
        self.channel.sendto(msg_g1, peer_addr)

        msg_u1 = self.channel.recv(1024)
        id_u, R_u, mac_u = unpack_bytes(msg_u1)
        if id_u != peer_id:
            logger.error("recv msg from unexpected peer [%s]", id_u.hex())
            # 在实际部署的实现中，这里应该继续等待合法消息，直到超时而不是直接终止
            # 此处为了实验方便设置为直接终止
            return

        if hmac.new(vk, id_u + R_u, digestmod="sha256").digest() != mac_u:
            # 在实际部署的实现中，这里应该继续等待合法消息，直到超时而不是直接终止
            # 此处为了实验方便设置为直接终止
            logger.error("msg with invalid mac from peer")
            return
        
        R_u = self.bytes_to_point(R_u)
        point_tmpkey = r_s * R_u
        tmpkey = hashlib.md5(self.point_to_bytes(point_tmpkey)).digest()
        #logger.debug("tmpkey is [%s]"%(tmpkey.hex()))
        ekey = AES.new(tmpkey, AES.MODE_ECB).encrypt(key)
        mac_g2 = hmac.new(tmpkey, id + ekey + ktype, digestmod="sha256").digest()
        # 封装第二条消息
        msg_g2 = pack_bytes([id, ekey, ktype, mac_g2])
        self.channel.sendto(msg_g2, peer_addr)

        msg_u2 = self.channel.recv(1024)
        id_u2, mac_u2 = unpack_bytes(msg_u2)
        if id_u2 != peer_id:
            logger.error("recv msg from unexpected peer [%s]", id_u.hex())
            # 在实际部署的实现中，这里应该继续等待合法消息，直到超时而不是直接终止
            # 此处为了实验方便设置为直接终止
            return
        if hmac.new(tmpkey, id_u2 + key, digestmod="sha256").digest() != mac_u2:
            logger.error("recv msg from unexpected peer [%s]", id_u.hex())
            # 在实际部署的实现中，这里应该继续等待合法消息，直到超时而不是直接终止
            # 此处为了实验方便设置为直接终止
            return
        logger.info("Success to set key [%s], ktype [%d]"%(key.hex(), int.from_bytes(ktype, "little")))
        end_time = time.time()
        logger.info("Protocol finished in [%.2f]ms"%((end_time - start_time) * 1000))

    # 根据id获取对应的公钥
    def get_pk(self, id) -> Point:
        return self.public_keys[id]

    def get_addr(self, id) -> tuple:
        return self.addrs[id]
    
    # ecpy使用的点转换函数
    # def point_to_bytes(self, point):
    #     return self.int_to_bytes(point.x) + self.int_to_bytes(point.y)
    # def bytes_to_point(self, data):
    #     size = int(self.curve.size / 8)
    #     x = int.from_bytes(data[:size], "big")
    #     y = int.from_bytes(data[size:], "big")
    #     return Point(x, y, self.curve)

    # pycryptodome使用的点转换函数
    def point_to_bytes(self, point):
        return point.x.to_bytes() + point.y.to_bytes()
    
    def bytes_to_point(self, data):
        size = int(self.curve.size / 8)
        x = Integer.from_bytes(data[:size])
        y = Integer.from_bytes(data[size:])
        return ECC.EccPoint(x, y)


    def int_to_bytes(self, num) -> bytes:
        return int.to_bytes(num, int(self.curve.size / 8), "big")


# basic implementation of dronesec lite protocol
class DronesecLite:
    def __init__(self, id : int, addrs, sk=None, keylist = None):
        self.id = id
        self.addrs = addrs
        # 长期密钥（预设参数）
        self.sk = sk
        self.keylist = keylist

        self.channel = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.channel.settimeout(4)
        self.channel.bind(addrs[id])
    
    def accept_key(self, station_id):
        if self.sk == None:
            logger.error("mode [gs] can't accept key")
            return
        id = int.to_bytes(self.id, 1, "little")
        station_id = int.to_bytes(station_id, 1, "little")
        logger.info("Protocol [dronesec_lite], waiting for key...")
        msg_g1, addr = self.channel.recvfrom(1024)

        start_time = time.time()

        id_g, ekey, mac_g1 = unpack_bytes(msg_g1)
        if station_id != id_g:
            logger.error("msg from unknown peer")
            return

        if hmac.new(self.sk, id_g + ekey, digestmod="sha256").digest() != mac_g1:
            logger.error("msg with invalid mac_g1")
            return
        
        key = AES.new(self.sk, AES.MODE_ECB).decrypt(ekey)
        ch_u = rand_bytes(16)
        mac_u1 = hmac.new(self.sk, id + ch_u + key, digestmod="sha256").digest()
        msg_u1 = pack_bytes([id, ch_u, mac_u1])
        self.channel.sendto(msg_u1, addr)

        msg_g2, addr = self.channel.recvfrom(1024)
        id_g2, ktype, mac_g2 = unpack_bytes(msg_g2)
        if id_g2 != id_g:
            logger.error("msg from unknown peer")
            return
        
        if hmac.new(key, id_g2 + ch_u + ktype, digestmod="sha256").digest() != mac_g2:
            logger.error("msg with invalid mac_g2")
            return
        
        mac_u2 = hmac.new(key, id + ch_u, digestmod="sha256").digest()
        msg_u2 = pack_bytes([id, mac_u2])
        self.channel.sendto(msg_u2, addr)
        ktype = int.from_bytes(ktype, "little")
        logger.info("success to get key [%s], ktype [%d]"%(key.hex(), ktype))
        end_time = time.time()
        logger.info("Protocol finished in [%.2f]ms"%((end_time - start_time) * 1000))
        return key, ktype

    def send_key(self, peer_id, key, ktype):
        # 读取对应无人机的长期密钥
        sk = self.get_longterm_key(peer_id)
        if sk == None:
            logger.error("mode [uav] can't send key")
            return
        # 读取无人机的ip地址和端口号
        peer_addr = self.get_addr(peer_id)
        #logger.debug("start to set key:\nprotocol [dronesec_lite]\nid [%d]\nkey [%s]\nktype [%s]"%(peer_id, key.hex(), ktype) )
        peer_id = int.to_bytes(peer_id, 1, "little")
        id = int.to_bytes(self.id, 1, "little")
        ktype = int.to_bytes(ktype, 1, "little")

        start_time = time.time()

        ekey = AES.new(sk, AES.MODE_ECB).encrypt(key)
        mac_g = hmac.new(sk, id + ekey, digestmod="sha256").digest()
        # 封装第一条消息
        msg_g1 = pack_bytes([id, ekey, mac_g])
        self.channel.sendto(msg_g1, peer_addr)

        msg_u1 = self.channel.recv(1024)
        id_u, ch_u, mac_u = unpack_bytes(msg_u1)
        if id_u != peer_id:
            logger.error("recv msg from unexpected peer [%s]", id_u.hex())
            # 在实际部署的实现中，这里应该继续等待合法消息，直到超时而不是直接终止
            # 此处为了实验方便设置为直接终止
            return

        if hmac.new(sk, id_u + ch_u + key, digestmod="sha256").digest() != mac_u:
            # 在实际部署的实现中，这里应该继续等待合法消息，直到超时而不是直接终止
            # 此处为了实验方便设置为直接终止
            logger.error("msg with invalid mac from peer")
            return
        
        mac_g2 = hmac.new(key, id + ch_u + ktype, digestmod="sha256").digest()
        # 封装第二条消息
        msg_g2 = pack_bytes([id, ktype, mac_g2])
        self.channel.sendto(msg_g2, peer_addr)

        msg_u2 = self.channel.recv(1024)
        id_u2, mac_u2 = unpack_bytes(msg_u2)
        if id_u2 != peer_id:
            logger.error("recv msg from unexpected peer [%s]", id_u.hex())
            # 在实际部署的实现中，这里应该继续等待合法消息，直到超时而不是直接终止
            # 此处为了实验方便设置为直接终止
            return
        if hmac.new(key, id_u2 + ch_u, digestmod="sha256").digest() != mac_u2:
            logger.error("recv msg from unexpected peer [%s]", id_u.hex())
            # 在实际部署的实现中，这里应该继续等待合法消息，直到超时而不是直接终止
            # 此处为了实验方便设置为直接终止
            return
        logger.info("Success to set key [%s], ktype [%d]"%(key.hex(), int.from_bytes(ktype, "little")))

        end_time = time.time()
        logger.info("Protocol finished in [%.2f]ms"%((end_time - start_time) * 1000))

    # 根据id获取对应的公钥
    def get_longterm_key(self, id):
        if self.keylist == None:
            return None
        else:
            return self.keylist[id]

    def get_addr(self, id) -> tuple:
        return self.addrs[id]

    def int_to_bytes(self, num) -> bytes:
        return int.to_bytes(num, int(self.curve.size / 8), "big")


# 生成随机私钥(随机大整数)
def generate_random(modulo):
    return random.randint(modulo >> 1, modulo)

def write_key_to_file(id, key, ktype):
    if ktype == 0:
        key_filename = "unicastkey"
    else:
        key_filename = "broadcastkey"
    with open(key_filename, "wb") as f:
        f.write(int.to_bytes(id, 1, "little"))
        f.write(key)
    logger.info("write key to [%s]"%(key_filename))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("mode", nargs=1, help="work mode, can be [gs] or [uav]")
    parser.add_argument("--id", default=0)
    parser.add_argument("--peer_id", default=1)
    parser.add_argument("--ktype","-t", default=0, help="0 for unicast key, 1 for broadcast key")
    parser.add_argument("--pswd", "-p", default="12345678", help="password for generating key")
    parser.add_argument("--protocol", default="dronesec", help="dronesec | dronsec_lite")
    args = parser.parse_args()
    # 读取命令行参数
    mode = args.mode[0]
    id = int(args.id)
    peer_id = int(args.peer_id)
    key = PBKDF2(args.pswd, b'qpalxntfcbg', dkLen=32, count=10000)

    if args.protocol == "dronesec_lite":
        if mode == "gs":
            protocol = DronesecLite(id, addrs, keylist=keys)
        else:
            protocol = DronesecLite(id, addrs, sk=keys[id])
    else:
        protocol = Dronesec(id, curve, private_keys[id], public_keys, addrs)
    if mode  == "gs":
        ktype = int(args.ktype)
        protocol.send_key(peer_id, key, ktype)
    else:
        while True:
            key, ktype = protocol.accept_key(peer_id)
            write_key_to_file(id, key, ktype)
            c = input("Key set finished, continue or quit?(default c): ")
            if c.startswith('q'):
                return

if __name__ == "__main__":
    main()
